import rclpy
from rclpy.node import Node
import whisper
import zhconv
import wave  # 使用wave库可读、写wav类型的音频文件
import pyaudio  # 使用pyaudio库可以进行录音，播放，生成wav文件
import threading


class SpeechDetect(Node):
    def __init__(self, name):
        super().__init__(name)
        # 定义数据流块
        self._CHUNK = 1024  # 音频帧率（也就是每次读取的数据是多少，默认1024）
        self._FORMAT = pyaudio.paInt16  # 采样时生成wav文件正常格式
        self._CHANNELS = 1  # 音轨数（每条音轨定义了该条音轨的属性,如音轨的音色、音色库、通道数、输入/输出端口、音量等。可以多个音轨，不唯一）
        self._RATE = 16000  # 采样率（即每秒采样多少数据）
        self._RECORD_SECONDS = 3  # 录音时间
        self._WAVE_OUTPUT_FILENAME = "./output.wav"  # 保存音频路径
        self._pyaudio = pyaudio.PyAudio()  # 创建PyAudio对象
        self._model = whisper.load_model("base")
        self._thd = threading.Thread(target=self._run)
        self._thd.start()

    def _record(self):
        stream = self._pyaudio.open(
            format=self._FORMAT,  # 采样生成wav文件的正常格式
            channels=self._CHANNELS,  # 音轨数
            rate=self._RATE,  # 采样率
            input=True,  # Ture代表这是一条输入流，False代表这不是输入流
            frames_per_buffer=self._CHUNK,
        )  # 每个缓冲多少帧
        self.get_logger().info("* recording")  # 开始录音标志
        frames = []  # 定义frames为一个空列表
        for i in range(
            0, int(self._RATE / self._CHUNK * self._RECORD_SECONDS)
        ):  # 计算要读多少次，每秒的采样率/每次读多少数据*录音时间=需要读多少次
            data = stream.read(self._CHUNK)  # 每次读chunk个数据
            frames.append(data)  # 将读出的数据保存到列表中
        self.get_logger().info("* done recording")  # 结束录音标志

        stream.stop_stream()  # 停止输入流
        stream.close()  # 关闭输入流
        # p.terminate()  # 终止pyaudio

        wf = wave.open(
            self._WAVE_OUTPUT_FILENAME, "wb"
        )  # 以’wb‘二进制流写的方式打开一个文件
        wf.setnchannels(self._CHANNELS)  # 设置音轨数
        wf.setsampwidth(
            self._pyaudio.get_sample_size(self._FORMAT)
        )  # 设置采样点数据的格式，和FOMART保持一致
        wf.setframerate(self._RATE)  # 设置采样率与RATE要一致
        wf.writeframes(b"".join(frames))  # 将声音数据写入文件
        wf.close()  # 数据流保存完，关闭文件

    def _recognize(self):
        result = self._model.transcribe("output.wav", language="Chinese")
        s = result["text"]
        s1 = zhconv.convert(s, "zh-cn")
        print(s1)

    def _run(self):
        while rclpy.ok():
            self._record()
            self._recognize()


def main():
    rclpy.init()
    node = SpeechDetect("speech")
    rclpy.spin(node)
    rclpy.shutdown()
